import math


def euclidean_distance(x1, x2):
    """
    计算两个数据点之间的欧氏距离
    
    参数:
        x1, x2: 数据点（列表或元组）
    
    返回:
        欧氏距离
    """
    return math.sqrt(sum((x1[i] - x2[i]) ** 2 for i in range(len(x1))))


def find_k_nearest_neighbors(x, X_train, k):
    """
    找到数据点x的k个最近邻
    
    参数:
        x: 待查询的数据点
        X_train: 训练数据点列表
        k: 近邻数量
    
    返回:
        k个最近邻的索引列表
    """
    distances = [(i, euclidean_distance(x, X_train[i])) for i in range(len(X_train))]
    distances.sort(key=lambda item: item[1])
    return [idx for idx, _ in distances[:k]]


def KNN_predict(x, X_train, y_train, k):
    """
    K近邻分类预测
    
    参数:
        x: 待预测的数据点
        X_train: 训练数据点列表
        y_train: 训练标签列表
        k: 近邻数量
    
    返回:
        预测的类别
    """
    neighbors_idx = find_k_nearest_neighbors(x, X_train, k)
    neighbor_labels = [y_train[i] for i in neighbors_idx]
    
    # 投票决定类别
    label_counts = {}
    for label in neighbor_labels:
        label_counts[label] = label_counts.get(label, 0) + 1
    
    # 返回出现次数最多的类别
    return max(label_counts, key=label_counts.get)


def KNN_predict_regression(x, X_train, y_train, k):
    """
    K近邻回归预测
    
    参数:
        x: 待预测的数据点
        X_train: 训练数据点列表
        y_train: 训练值列表
        k: 近邻数量
    
    返回:
        预测值（k个近邻的平均值）
    """
    neighbors_idx = find_k_nearest_neighbors(x, X_train, k)
    neighbor_values = [y_train[i] for i in neighbors_idx]
    return sum(neighbor_values) / len(neighbor_values)


# 测试示例
if __name__ == "__main__":
    # 分类示例
    print("=== KNN分类示例 ===")
    X_train = [
        [1, 1], [1, 2], [2, 1], [2, 2],
        [5, 5], [5, 6], [6, 5], [6, 6],
        [9, 9], [9, 10], [10, 9], [10, 10]
    ]
    y_train = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]
    
    test_point = [3, 3]
    k = 3
    prediction = KNN_predict(test_point, X_train, y_train, k)
    print(f"测试点 {test_point} 的预测类别: {prediction}")
    
    # 回归示例
    print("\n=== KNN回归示例 ===")
    X_train_reg = [[1], [2], [3], [4], [5], [6], [7], [8], [9], [10]]
    y_train_reg = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
    
    test_point_reg = [5.5]
    k_reg = 3
    prediction_reg = KNN_predict_regression(test_point_reg, X_train_reg, y_train_reg, k_reg)
    print(f"测试点 {test_point_reg} 的预测值: {prediction_reg}")